library(GenomicRanges)
RleSumAny <- function (e1, e2)
{
    len <- length(e1)
    stopifnot(len == length(e2))
    x1 <- runValue(e1); s1 <- cumsum(runLength(e1))
    x2 <- runValue(e2); s2 <- cumsum(runLength(e2))
    .Call("rle_sum_any",
          as.integer(x1), as.integer(s1),
          as.integer(x2), as.integer(s2),
          as.integer(len),
          PACKAGE = "chipseq")
}


GetGRanges <- function(LoadFile,AllChr=NULL,ChrOfInterest=NULL,simple=F,sepr="\t",simplify=F){
   require(Rsamtools)
   require(GenomicRanges)

   if(class(LoadFile) == "GRanges"){
      RegionRanges <- LoadFile
       if(simplify){
         RegionRanges <- GRanges(seqnames(RegionRanges),ranges(RegionRanges))
       }
   }else{
    if(class(LoadFile) == "character"){
      RangesTable <- read.delim(LoadFile,sep=sepr,header=T,skip="#")
    }else if(class(LoadFile) == "matrix"){
      RangesTable <- as.data.frame(LoadFile)
    } else{
      RangesTable <- as.data.frame(LoadFile)
    }
    Chromosomes <- as.vector(RangesTable[,1])
    Start <- as.numeric(as.vector(RangesTable[,2]))
    End <- as.numeric(as.vector(RangesTable[,3]))
    RegionRanges <- GRanges(seqnames=Chromosomes,ranges=IRanges(start=Start,end=End))
    if(simple == F){
    if(ncol(RangesTable) > 4){
      ID <- as.vector(RangesTable[,4])
      Score <- as.vector(RangesTable[,5])
      if(ncol(RangesTable) > 6){
        Strand <- rep("*",nrow(RangesTable))
        RemainderColumn <- as.data.frame(RangesTable[,-c(1:6)])
        elementMetadata(RegionRanges) <- cbind(ID,Score,Strand,RemainderColumn)
      }else{
        elementMetadata(RegionRanges) <- cbind(ID,Score)
      }
    }
    }
    }
    if(!is.null(AllChr)){ 
      RegionRanges[seqnames(RegionRanges) %in% AllChr]    
      seqlevels(RegionRanges,force=T) <- AllChr
    }
    if(!is.null(ChrOfInterest)){      
      RegionRanges <- RegionRanges[seqnames(RegionRanges) == ChrOfInterest]      
    }

    return(RegionRanges)
}



findCovMaxPos <- function(reads,bedRanges,ChrOfInterest,FragmentLength){
  require(GenomicRanges)
  require(Rsamtools)

      cat("done\n")
      cat("Calculating coverage\n")
      MaxRanges <- GRanges()
      if(length(reads) > 0){
      AllCov <- coverage(reads)
      AllCov <- runmean(AllCov[names(AllCov) %in% ChrOfInterest],20)
      cat("Calculating Summits on ",ChrOfInterest," ..")
      Maxes <- which.max(Views(AllCov[[which(names(AllCov) %in% ChrOfInterest)]],ranges(bedRanges[seqnames(bedRanges) == ChrOfInterest])))
      if(any(is.na(Maxes))){ 
        NoSummitRanges <- bedRanges[is.na(Maxes)]
        Maxes[is.na(Maxes)]  <- (start((ranges(NoSummitRanges[seqnames(NoSummitRanges) == ChrOfInterest])))+end((ranges(NoSummitRanges[seqnames(NoSummitRanges) == ChrOfInterest]))))/2
      }
      MaxRanges <- GRanges(seqnames(bedRanges[seqnames(bedRanges) == ChrOfInterest]),IRanges(start=Maxes,end=Maxes),elementMetadata=elementMetadata(bedRanges[seqnames(bedRanges) == ChrOfInterest]))
      cat(".done\n")
      }
  return(MaxRanges)
}


setClass("ChIPQCsample",contains = "GRanges",
slots=c(AveragePeakSignal="list",
CrossCoverage="numeric",SSD="numeric",CountsInPeaks="numeric",
CountsInBlackList="numeric",CountsInFeatures="list",CoverageHistogram="numeric",FlagAndTagCounts="numeric"
))


setMethod("show","ChIPQCsample",
  function (object){
    message("\t\t\t\t\t",class(object),"")
    message("Number of Mapped reads: ",object@FlagAndTagCounts[2],"")
    message("Number of Mapped reads passing MapQ filter: ",object@FlagAndTagCounts[4],"")    
    message("Percentage Of Reads as Non-Duplicates (NRF): ",round((((object@FlagAndTagCounts[2])-(object@FlagAndTagCounts[5]))/object@FlagAndTagCounts[2])*100, digits=2),"(",round(object@FlagAndTagCounts[5]/object@FlagAndTagCounts[2],digits=2),")","")            
    message("Percentage Of Reads in Blacklisted Regions: ",round((object@CountsInBlackList/object@FlagAndTagCounts[2])*100),"")            

    message("SSD: ",object@SSD,"")
    message("Fragment Length Cross-Coverage: ",FragmentLengthCrossCoverage(object),"")
    message("Relative Cross-Coverage: ",RelativeCrossCoverage(object),"")
    
    message("Percentage Of Reads in GenomicFeature: ")
    print(data.frame(ProportionOfCounts=unlist(object@CountsInFeatures),row.names=gsub("CountsIn","",names(object@CountsInFeatures)))/object@FlagAndTagCounts[4])
   
    message("Percentage Of Reads in Peaks: ",round((object@CountsInPeaks/object@FlagAndTagCounts[4])*100,digits=2),"")            
    message("Number of Peaks: ",length(object),"")    

    print(as(object,"GRanges"))
  }
)

setGeneric("crosscoverage", function(object="ChIPQCsample") standardGeneric("crosscoverage"))
setMethod("crosscoverage", signature(object="ChIPQCsample"), function(object) ((object@CrossCoverage[1]-object@CrossCoverage)/object@CrossCoverage[1]))
setGeneric("ssd", function(object="ChIPQCsample") standardGeneric("ssd"))
setMethod("ssd", "ChIPQCsample", function(object) object@SSD)
setGeneric("fragmentlength", function(object="ChIPQCsample",width) standardGeneric("fragmentlength"))
setMethod("fragmentlength", "ChIPQCsample", function(object,width=3){
    MaxShift <- which.max(running(crosscoverage(object)[-c(1:(2*readlength))],width=width,allow=T))+2*readlength
    return(unname(MaxShift))
})

setGeneric("FragmentLengthCrossCoverage", function(object="ChIPQCsample",width) standardGeneric("FragmentLengthCrossCoverage"))
setMethod("FragmentLengthCrossCoverage", signature(object="ChIPQCsample"), function(object){ 
  FragmentLengthCrossCoverage <- crosscoverage(object)[fragmentlength(object,10)]-crosscoverage(object)[1]
  return(FragmentLengthCrossCoverage)
}
)
setGeneric("ReadLengthCrossCoverage", function(object="ChIPQCsample",width) standardGeneric("ReadLengthCrossCoverage"))
setMethod("ReadLengthCrossCoverage", signature(object="ChIPQCsample"), function(object){ 
  ReadLengthCrossCoverage <-  crosscoverage(object)[readlength]-crosscoverage(object)[1]
  return(ReadLengthCrossCoverage)
}
)

setGeneric("RelativeCrossCoverage", function(object="ChIPQCsample",width) standardGeneric("RelativeCrossCoverage"))
setMethod("RelativeCrossCoverage", signature(object="ChIPQCsample"), function(object){ 
  RelativeCrossCoverage <- FragmentLengthCrossCoverage(object)/ReadLengthCrossCoverage(object)
  return(RelativeCrossCoverage)
}
)


setGeneric("flagtagcounts", function(object="ChIPQCsample") standardGeneric("flagtagcounts"))
setMethod("flagtagcounts", "ChIPQCsample", function(object) object@FlagAndTagCounts)
setGeneric("coveragehistogram", function(object="ChIPQCsample") standardGeneric("coveragehistogram"))
setMethod("coveragehistogram", "ChIPQCsample", function(object) object@CoverageHistogram)
setGeneric("averagepeaksignal", function(object="ChIPQCsample") standardGeneric("averagepeaksignal"))
setMethod("averagepeaksignal", "ChIPQCsample", function(object) object@AveragePeakSignal[[1]])
setGeneric("Normalisedaveragepeaksignal", function(object="ChIPQCsample") standardGeneric("Normalisedaveragepeaksignal"))
setMethod("Normalisedaveragepeaksignal", "ChIPQCsample", function(object) object@AveragePeakSignal[[2]])
setGeneric("peaks", function(object="ChIPQCsample") standardGeneric("peaks"))
setMethod("peaks", "ChIPQCsample", function(object) as(object,"GRanges"))





ChIPQC <- function(bamFile,bedFile,blklist=NULL,ChrOfInterest=NULL,GeneAnnotation=NULL,Window=400,FragmentLength=50,shiftWindowStart=1,shiftWindowEnd=2,mapQCutoff=15){
    require(Rsamtools)
    require(GenomicRanges)
    require(GenomicAlignments)
    
    ChrLengths <- scanBamHeader(bamFile)[[1]]$targets
    message("Bam file has ",length(names(ChrLengths))," contigs")
    if(!all(ChrOfInterest %in% names(ChrLengths))){
      stop("Contigs of interest are not all in Bam file!!")
    }
    if(!is.null(ChrOfInterest)){
        ChrLengths <- ChrLengths[names(ChrLengths) %in% ChrOfInterest]
    }

    ShiftMat <- NULL
    FlagTagCounts <- NULL
    CovHist <- NULL
    SSD <- NULL
    CoverageMatrix <- NULL
    CoverageMatrixInput <- NULL
    bedRangesTemp <- GetGRanges(GRanges(),AllChr=names(ChrLengths))
    bedRangesSummitsTemp <- vector("numeric")
    Counts <- NULL
    MultiRangesCountsTemp <- NULL
    txdb <- NULL
    if(!is.null(GeneAnnotation)){
        if(GeneAnnotation == "hg19"){
          txdb <- TxDb.Hsapiens.UCSC.hg19.knownGene
        }
      All5utrs <- reduce(unique(unlist(fiveUTRsByTranscript(txdb))))
      All3utrs <- reduce(unique(unlist(threeUTRsByTranscript(txdb))))
      Allcds <- reduce(unique(unlist(cdsBy(txdb,"tx"))))
      Allintrons <- reduce(unique(unlist(intronsByTranscript(txdb))))
      Alltranscripts <- reduce(unique(unlist(transcripts(txdb))))
      Promoters500 <-  reduce(flank(Alltranscripts,500,"start"))    
      Promoters2000t0500 <-  reduce(flank(Promoters500,2000,"start"))
      LongPromoter20000to2000  <- reduce(flank(Promoters2000t0500,20000,"start"))
    
    }
    

    for(k in 1:length(ChrLengths)){

       Param <- ScanBamParam(which=GRanges(seqnames=names(ChrLengths)[k],IRanges(start=1,end=unname(ChrLengths[names(ChrLengths) == names(ChrLengths)[k]])-shiftWindowEnd)),
       what=c("flag","mapq"))
       temp <- readGAlignmentsFromBam(bamFile,param=Param)
    
    
       Sample_GIT <- GIntervalTree(GRanges(seqnames=seqnames(temp),ranges=ranges(temp),strand=strand(temp),elementMetadata(temp)))
       
       flagMapQ <- cbind(bamFlagAsBitMatrix(Sample_GIT$flag)[,c("isUnmappedQuery","isDuplicate")],Sample_GIT$mapq)
       colnames(flagMapQ) <- c("A","B","C")
       temp <- as.data.frame(xtabs(~A+B+C, data=flagMapQ))
       UnMapped <- sum(temp[temp[,"A"] == 1,"Freq"])
       Mapped <- sum(temp[temp[,"A"] != 1,"Freq"])
       Duplicates <- sum(temp[temp[,"A"] != 1 & temp[,"B"] == 1,"Freq"])
    
       MapQPass <- sum(temp[temp[,"A"] != 1 & as.numeric(temp[,"C"]) >= 15,"Freq"])
       MapQPassAndDup <- sum(temp[temp[,"A"] != 1 & temp[,"B"] == 1 & as.numeric(temp[,"C"]) >= mapQCutoff,"Freq"])
       FlagTagCounts <- rbind(cbind(UnMapped,Mapped,Duplicates,MapQPass,MapQPassAndDup),FlagTagCounts)

    
       Cov <- coverage(Sample_GIT)
       message("Calculating coverage histogram..\n")
       CovHist <- c(CovHist,list(colSums(table(Cov))))
       message("Calculating SSD..\n")
       SSD <- c(SSD,sd(Cov))
    
       PosCoverage <- coverage(IRanges(start(Sample_GIT[strand(Sample_GIT)=="+"]),start(Sample_GIT[strand(Sample_GIT)=="+"])),width=ChrLengths[k])
       NegCoverage <- coverage(IRanges(end(Sample_GIT[strand(Sample_GIT)=="-"]),end(Sample_GIT[strand(Sample_GIT)=="-"])),width=ChrLengths[k])
       message("Calculating shift for ",names(ChrLengths)[k],"\n")
       #ShiftsTemp <- shiftApply(seq(shiftWindowStart,shiftWindowEnd),PosCoverage,NegCoverage,cor)
       ShiftsTemp <- shiftApply(seq(shiftWindowStart,shiftWindowEnd),PosCoverage,NegCoverage,RleSumAny, verbose = T)
       ShiftMat <- cbind(ShiftMat,ShiftsTemp)
       
       bedRanges <- GetGRanges(bedFile,as.vector(names(ChrLengths)),names(ChrLengths)[k])
  

       GRangesOfInterestList <- GRangesList(GRanges(seqnames(bedRanges),ranges(bedRanges)))
       if(!is.null(blklist)){
         blkRanges <- GetGRanges(blklist,names(ChrLengths),names(ChrLengths)[k])           
         GRangesOfInterestList <- c(GRangesOfInterestList,GRangesList(GRanges(seqnames(blkRanges),ranges(blkRanges))))
       }
      if(!is.null(txdb)){
        GRangesOfInterestList <- c(GRangesOfInterestList,
        GRangesList(GetGRanges(All5utrs,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(All3utrs,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(Allcds,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(Allintrons,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(Alltranscripts,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(Promoters500,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(Promoters2000t0500,names(ChrLengths),names(ChrLengths)[k],simplify=T),
        GetGRanges(LongPromoter20000to2000,names(ChrLengths),names(ChrLengths)[k],simplify=T)))
      }
       MultiRangesCountsTemp <- c(MultiRangesCountsTemp,countOverlaps(GRangesOfInterestList,Sample_GIT))
       message("Counting reads in peaks..","\n")        
       CountsTemp <- countOverlaps(bedRanges,Sample_GIT)
       Counts  <- c(Counts,CountsTemp)

       bedRangesTemp <- c(bedRangesTemp,bedRanges)
       message("Signal over peaks..","\n")                   
       AllFragRanges <- resize(as(Sample_GIT[Sample_GIT %over% bedRanges],"GRanges"),FragmentLength,"start")
       bedRangesSummits <- findCovMaxPos(AllFragRanges,bedRanges,names(ChrLengths)[k],FragmentLength)
       bedRangesSummitsTemp <- c(bedRangesSummitsTemp,as.numeric(as.vector(start(bedRangesSummits))))
       updatedRanges  <- resize(bedRangesSummits,Window,"center")
       AllCov <- coverage(AllFragRanges,width=unname(ChrLengths[k])+Window)
       CoverageMatrix <- rbind(CoverageMatrix,matrix(as.vector(AllCov[[which(names(AllCov) == names(ChrLengths)[k])]][ranges(updatedRanges[seqnames(updatedRanges) == names(ChrLengths)[k]])]),ncol=Window,byrow=T))

    }

    FlagTagCounts <- colSums(FlagTagCounts)
    AvProfile <- colMeans(CoverageMatrix)
    NormAvProfile <- (AvProfile/FlagTagCounts[4])*1e6
    Weights <- ChrLengths
    CovHistAll <- NULL
    for(l in CovHist){
      CovHistAll <- rowSums(merge(CovHistAll,cbind(as.numeric(names(l)),as.numeric(l)),by.x=0,by.y=1,all=T)[,-1,drop=F])

    }
      
    elementMetadata(bedRangesTemp) <- data.frame(Counts,bedRangesSummitsTemp)
    print(length(GRangesOfInterestList))
    ReadsInPeaks <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,1])
    ReadsInBLs <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,2])  
    ReadsIn5utrs<- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,3])
    ReadsIn3utrs <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,4])
    ReadsInCDS <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,5])
    ReadsInIntrons <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,6])
    ReadsInTranscripts <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,7])              
    ReadsInPromoters500 <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,8])              
    ReadsInPromoters2000t0500 <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,9])              
    ReadsInLongPromoter20000to2000 <- sum(matrix(MultiRangesCountsTemp,ncol=length(GRangesOfInterestList),byrow=T)[,10])              
    
    
    
    
    
    ShiftsAv <- apply(ShiftMat,1,function(x)weighted.mean(x,Weights,na.rm=T))
    SSDAv <- (weighted.mean(sd(Cov)[names(ChrLengths)],ChrLengths)*1000)/sqrt(FlagTagCounts[4])
    CHQC <- new("ChIPQCsample", bedRangesTemp,AveragePeakSignal=list(AvProfile,NormAvProfile),
    CrossCoverage=ShiftsAv,SSD=SSDAv,CountsInPeaks=ReadsInPeaks,CountsInBlackList=ReadsInBLs,
    CountsInFeatures=list(CountsIn5utrs=ReadsIn5utrs,CountsIn3UTRs=ReadsIn3utrs,
    CountsinReads=ReadsInCDS,CountsInIntrons=ReadsInIntrons,CountsInTranscripts=ReadsInTranscripts,
    CountsInPromoters500=ReadsInPromoters500,CountsInPromoters2000t0500=ReadsInPromoters2000t0500,CountsInLongPromoter20000to2000=ReadsInLongPromoter20000to2000),
    CoverageHistogram=CovHistAll,FlagAndTagCounts=FlagTagCounts)
#    elementMetadata(CHQC) <- Counts    
    print(MultiRangesCountsTemp)
    return(CHQC)    

}



